package com.github.netty.util;

import javax.servlet.annotation.WebServlet;
import java.io.File;
import java.io.IOException;
import java.lang.annotation.Annotation;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.JarURLConnection;
import java.net.URL;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

public class SecureAnnotationScanner {

    public static Set<Class<?>> scanPackageForAnnotations(String basePackage, Class<? extends Annotation> annotationClass) {
        Set<Class<?>> classesWithAnnotation = new HashSet<>();
        URL url = SecureAnnotationScanner.class.getResource("/" + basePackage.replace('.', '/'));

        if (url != null) {
            try {
                if ("file".equals(url.getProtocol())) {
                    File directory = new File(url.toURI());
                    classesWithAnnotation.addAll(scanDirectory(directory, annotationClass));
                } else if ("jar".equals(url.getProtocol())) {
                    JarURLConnection connection = (JarURLConnection) url.openConnection();
                    JarFile jarFile = connection.getJarFile();
                    classesWithAnnotation.addAll(scanJarFile(jarFile, basePackage, annotationClass));
                }
            } catch (Exception e) {
                throw new RuntimeException("Error scanning package for annotations", e);
            }
        }
        return classesWithAnnotation;
    }

    private static Set<Class<?>> scanDirectory(File directory, Class<? extends Annotation> annotationClass) throws IOException, ClassNotFoundException {
        Set<Class<?>> annotatedClasses = new HashSet<>();
        Files.walk(Paths.get(directory.toURI()))
             .filter(Files::isRegularFile)
             .filter(p -> p.toString().endsWith(".class"))
             .forEach(filePath -> {
                 try {
                     String className = classNameFromPath(filePath.toString());
                     Class<?> clazz = Class.forName(className);
                     if (clazz.isAnnotationPresent(annotationClass)) {
                         annotatedClasses.add(clazz);
                     }
                 } catch (ClassNotFoundException e) {
                     throw new RuntimeException(e);
                 }
             });
        return annotatedClasses;
    }

    private static Set<Class<?>> scanJarFile(JarFile jarFile, String basePackage, Class<? extends Annotation> annotationClass) throws IOException {
        Set<Class<?>> annotatedClasses = new HashSet<>();
        Enumeration<JarEntry> entries = jarFile.entries();
        while (entries.hasMoreElements()) {
            JarEntry entry = entries.nextElement();
            if (entry.getName().startsWith(basePackage) && entry.getName().endsWith(".class") && !entry.getName().contains("$")) {
                String className = entry.getName().replace('/', '.').substring(0, entry.getName().length() - 6);
                try {
                    Class<?> clazz = Class.forName(className);
                    if (clazz.isAnnotationPresent(annotationClass)) {
                        annotatedClasses.add(clazz);
                    }
                } catch (ClassNotFoundException e) {
                    throw new RuntimeException(e);
                }
            }
        }
        return annotatedClasses;
    }

    private static String classNameFromPath(String path) {
        return path.replace(File.separator, ".").replace("\\.class", "");
    }

    public static void main(String[] args) {
        Set<Class<?>> annotatedClasses = scanPackageForAnnotations("com.example", WebServlet.class);
        for (Class<?> clazz : annotatedClasses) {
            System.out.println("Class with @WebServlet: " + clazz.getName());
        }
    }
}